import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import time
import torch
from torch import nn, optim
from torch.utils import data
from dataloder import CarDateSet
import argparse
from torchvision.models import resnet50, resnext50_32x4d, densenet121
from torchvision.models import *

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device:', device)


def main(args):
    # Create model
    if not os.path.exists(args.model_path):
        os.makedirs(args.model_path)

    train_datasets = CarDateSet(args.dataset_path + 'train_data', args.dataset_path + 'train.txt', transforms=None)
    valid_datasets = CarDateSet(args.dataset_path + 'test_data', args.dataset_path + 'test.txt', transforms=None)

    train_loader = torch.utils.data.DataLoader(dataset=train_datasets,
                                               batch_size=args.batch_size,
                                               shuffle=True)

    test_loader = torch.utils.data.DataLoader(dataset=valid_datasets,
                                              batch_size=args.batch_size,
                                              shuffle=True)

    print("Train numbers:{:d}".format(len(train_datasets)))
    print("Test numbers:{:d}".format(len(valid_datasets)))

    model = eval(args.net)(pretrained=args.pretrained)
    # model = resnext50_32x4d(pretrained=args.pretrained)
    if not args.pretrained:
        model.load_state_dict(torch.load(args.pretrained_model))  # 加载模型参数
    channel_in_2 = model.fc.in_features  # 获取fc层的输入通道数
    # 然后把resnet的fc层替换成自己分类类别的fc层
    model.fc = nn.Linear(channel_in_2, args.num_class)
    print('model parameters:', sum(p.numel() for p in model.parameters() if p.requires_grad))
    model = model.to(device)
    cost2 = nn.CrossEntropyLoss().to(device)
    optimizer2 = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-6)
    best_acc_2 = 0.

    for epoch in range(1, args.epochs + 1):
        model.train()
        # start time
        start = time.time()
        index = 0
        for images, labels in train_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs2 = model(images)
            loss2 = cost2(outputs2, labels)
            optimizer2.zero_grad()
            loss2.backward()
            optimizer2.step()
            index += 1

        if epoch % 1 == 0:
            end = time.time()
            print("Epoch [%d/%d], Loss: %.8f, Time: %.1fsec!" % (epoch, args.epochs, loss2.item(), (end - start) * 2))
            model.eval()
            class_correct2 = list(0. for i in range(args.num_class))
            class_total2 = list(0. for i in range(args.num_class))
            correct_prediction_2 = 0.
            total_2 = 0
            with torch.no_grad():
                for images, labels in test_loader:
                    # to GPU
                    images = images.to(device)
                    labels = labels.to(device)
                    outputs2 = model(images)
                    _2, predicted2 = torch.max(outputs2, 1)
                    c2 = (predicted2 == labels).squeeze()
                    for label_idx in range(len(labels)):
                        label = labels[label_idx]
                        class_correct2[label] += c2[label_idx].item()
                        class_total2[label] += 1
                    total_2 += labels.size(0)
                    # add correct
                    correct_prediction_2 += (predicted2 == labels).sum().item()

            acc_2 = correct_prediction_2 / total_2
            print('----------------------------------------------------')
        if acc_2 > best_acc_2:
            print('save new best acc_2', acc_2)
            torch.save(model, os.path.join(args.model_path, '{}_{}.pth'.format(args.net, args.dataset_name)))
            best_acc_2 = acc_2
    print("Model save to %s." % (os.path.join(args.model_path, '{}_{}.pth'.format(args.net, args.dataset_name))))
    print('save new best acc_2', best_acc_2)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='train hyper-parameter')
    parser.add_argument("--num_class", default=43, type=int)
    parser.add_argument("--epochs", default=100, type=int)
    parser.add_argument("--net", default='ResNet50', type=str)
    parser.add_argument("--lr", default=1e-5, type=float)
    parser.add_argument("--batch_size", default=32, type=int)
    parser.add_argument("--dataset_name", default='GID', type=str)
    parser.add_argument("--dataset_path", default='../GID/', type=str)
    parser.add_argument("--model_path", default='SaveModel', type=str)
    parser.add_argument("--pretrained", default=False, type=bool)
    parser.add_argument("--pretrained_model", default='PretrainedModel/resnext50_32x4d-7cdf4587.pth', type=str)
    args = parser.parse_args()

    main(args)
